package security;

import java.io.IOException;
import java.security.InvalidAlgorithmParameterException;
import java.security.InvalidKeyException;
import java.security.NoSuchAlgorithmException;
import java.security.spec.AlgorithmParameterSpec;

import javax.crypto.BadPaddingException;
import javax.crypto.Cipher;
import javax.crypto.IllegalBlockSizeException;
import javax.crypto.NoSuchPaddingException;
import javax.crypto.SecretKey;

import org.bouncycastle.util.encoders.Base64;

/**
 * Provides AES-encrypted communication.
 * 
 * @author Lukas
 *
 */
public class AESChannel extends ChannelDecorator {
	private Cipher encrypt, decrypt;
	
	/**
	 * Constructor.
	 * @param channel The underlying channel that shall be used for transmission.
	 * @param secKey The secret key.
	 * @param ivPar The iv-parameter.
	 * @throws IOException
	 */
	public AESChannel(Channel channel, SecretKey secKey, AlgorithmParameterSpec ivPar) throws IOException {
		this.channel = channel;
		initCiphers(secKey, ivPar);
	}

	@Override
	public void send(String message) throws IOException {
		try {
			channel.send(new String(Base64.encode(encrypt.doFinal(message.getBytes()))));
		} catch (IllegalBlockSizeException e) {
			throw new IOException(e);
		} catch (BadPaddingException e) {
			throw new IOException(e);
		}
	}

	@Override
	public String receive() throws IOException {
		String message = channel.receive();
		if (message.equals("!unsecure") || message.equals("!lastSecure")) {
			return message;
		}
		try {
			return new String(decrypt.doFinal(Base64.decode(message.getBytes())));
		} catch (IllegalBlockSizeException e) {
			throw new IOException(e);
		} catch (BadPaddingException e) {
			throw new IOException(e);
		}
	}

	private void initCiphers(SecretKey secKey, AlgorithmParameterSpec ivPar) throws IOException {
		try {
			encrypt = Cipher.getInstance("AES/CTR/NoPadding");
			decrypt = Cipher.getInstance("AES/CTR/NoPadding");
		} catch (NoSuchAlgorithmException e) {
			throw new IOException(e);
		} catch (NoSuchPaddingException e) {
			throw new IOException(e);
		}
		try {
			encrypt.init(Cipher.ENCRYPT_MODE, secKey, ivPar);
			decrypt.init(Cipher.DECRYPT_MODE, secKey, ivPar);
		} catch (InvalidKeyException e) {
			throw new IOException(e);
		} catch (InvalidAlgorithmParameterException e) {
			throw new IOException(e);
		}
	}
}
